import time

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response
from sqlalchemy.orm import Session
from models.logs import RequestLog
from database import SessionLocal
from typing import Callable

from utils.token_utils import parse_token
from utils.log_utils import logger


class RequestLoggerMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: Callable):
        # 获取请求体
        try:
            body_bytes = await request.body()
            request_body = body_bytes.decode("utf-8", errors="ignore")[:1000]

            async def receive():
                return {"type": "http.request", "body": body_bytes}
            request._receive = receive
        except Exception:
            request_body = ""

        # 记录开始时间
        start = time.time()

        # 调用实际业务
        try:
            response = await call_next(request)
        except Exception as e:
            logger.error(f"中间件执行失败: {e}")
            return Response("Internal Server Error", status_code=500)

        # 响应体处理
        try:
            response_body = b""
            async for chunk in response.body_iterator:
                response_body += chunk
            response_text = response_body.decode("utf-8", errors="ignore")[:1000]

            async def new_body_iterator():
                yield response_body
            response.body_iterator = new_body_iterator()
        except Exception:
            response_text = ""

        # 响应耗时写入 header
        duration = round(time.time() - start, 5)
        response.headers["X-Process-Time"] = str(duration)

        # 用户身份解析
        try:
            auth_header = request.headers.get("Authorization")
            user = parse_token(auth_header[7:]) if auth_header else None
            user_id = user.get("id")
            username = user.get("sub")
        except Exception:
            user_id = None
            username = None

        # IP、UA 等信息
        ip = request.headers.get("x-forwarded-for", request.client.host)
        user_agent = request.headers.get("user-agent", "")

        # 日志落库
        try:
            db: Session = SessionLocal()
            log = RequestLog(
                user_id=user_id,
                username=username,
                method=request.method,
                url=str(request.url),
                status_code=response.status_code,
                client_ip=ip,
                user_agent=user_agent,
                request_body=request_body,
                response_body=response_text,
            )
            db.add(log)
            db.commit()
        except Exception as e:
            db.rollback()
            logger.error(f"日志保存失败: {e}")
        finally:
            db.close()

        return response
